import json
import logging
from channels.generic.websocket import WebsocketConsumer
from zhipuai import ZhipuAI

class ChatConsumer(WebsocketConsumer):
    client = ZhipuAI(api_key="1623ba7fdb8373a4db64addcdb52d9e9.ilwdshbLnW5FhsYg")
    def connect(self):
        self.accept()
        logging.info('connect')

    def disconnect(self, close_code):
        logging.info(f'disconnect:{close_code}')
        pass

    def receive(self, text_data):
        logging.info(f"recv:{text_data}")
        text_data_json = json.loads(text_data)
        message = text_data_json['info']
        self.__getAnswer(message)

    def __getAnswer(self, info: str):
        response = self.client.chat.completions.create(
            model="glm-4",  # 填写需要调用的模型名称
            messages=[
                {"role": "system", "content": "熟悉各种xml标准的语法，精准为用户提供帮助"},
                {"role": "user", "content": f"内容如下:\n{info}\n转换要求:将以上内容转换为dita标准的xml；中文显示；只输出xml内容。"},
            ],
            stream=True,
        )
        result = ''
        for chunk in response:
            # print(chunk.choices[0].delta.content)
            self.send(text_data=chunk.choices[0].delta.content)
            result += chunk.choices[0].delta.content
        # self.send(text_data=json.dumps({
        #     'flag': 'over'
        # }))
        logging.info(f"result:{result}")